#include "ant_colony_optimization.h"
#include <cmath>
#include <limits>
#include <ctime>
#include <cstdlib>
#include <algorithm>
#include <QDebug>

// 定义无穷大的值
const double INF = std::numeric_limits<double>::max();

// 构造函数
AntColonyOptimization::AntColonyOptimization(const QVector<QVector<double>>& distanceMatrix, int numAnts, int generations, double alpha, double beta, double evaporationRate)
    : distanceMatrix(distanceMatrix), numAnts(numAnts), generations(generations), alpha(alpha), beta(beta), evaporationRate(evaporationRate) {
    numberOfCities = distanceMatrix.size();
    srand(static_cast<unsigned>(time(0))); // 用当前时间作为随机种子
    pheromoneMatrix.resize(numberOfCities, QVector<double>(numberOfCities, 1.0)); // 初始化信息素矩阵为1.0
}

// 设置和获取函数
void AntColonyOptimization::setNumAnts(int numAnts) { this->numAnts = numAnts; }// 设置蚂蚁数量
int AntColonyOptimization::getNumAnts() const { return numAnts; }// 获取蚂蚁数量
void AntColonyOptimization::setGenerations(int generations) { this->generations = generations; }// 设置迭代代数
int AntColonyOptimization::getGenerations() const { return generations; }// 获取迭代代数
void AntColonyOptimization::setAlpha(double alpha) { this->alpha = alpha; }// 设置信息素重要性因子
double AntColonyOptimization::getAlpha() const { return alpha; }// 获取信息素重要性因子
void AntColonyOptimization::setBeta(double beta) { this->beta = beta; }// 设置启发式信息重要性因子
double AntColonyOptimization::getBeta() const { return beta; }// 获取启发式信息重要性因子
void AntColonyOptimization::setEvaporationRate(double rate) { evaporationRate = rate; }// 设置信息素蒸发率
double AntColonyOptimization::getEvaporationRate() const { return evaporationRate; }// 获取信息素蒸发率

// 计算路径长度
double AntColonyOptimization::calculateDistance(const QVector<int>& path) {
    double totalDistance = 0.0;
    for (int i = 0; i < numberOfCities - 1; ++i) {
        double distance = distanceMatrix[path[i]][path[i + 1]];
        if (distance == 0.0) distance = INF; // 将0视为无穷大
        totalDistance += distance;
    }
    totalDistance += distanceMatrix[path[numberOfCities - 1]][path[0]]; // 回到起点
    return totalDistance;
}

// 初始化信息素矩阵
void AntColonyOptimization::initializePheromones() {
    pheromoneMatrix = QVector<QVector<double>>(numberOfCities, QVector<double>(numberOfCities, 1.0)); // 初始化为1.0
}

// 更新信息素矩阵
void AntColonyOptimization::updatePheromones() {
    // 蒸发：所有路径的信息素强度按蒸发率减少
    for (int i = 0; i < numberOfCities; ++i) {
        for (int j = 0; j < numberOfCities; ++j) {
            pheromoneMatrix[i][j] *= (1.0 - evaporationRate);
        }
    }

    // 增加信息素：根据蚂蚁路径长度更新信息素
    for (int ant = 0; ant < numAnts; ++ant) {
        double pathLength = antPathLengths[ant];
        for (int i = 0; i < numberOfCities - 1; ++i) {
            int from = antPaths[ant][i];
            int to = antPaths[ant][i + 1];
            pheromoneMatrix[from][to] += 1.0 / pathLength;
            pheromoneMatrix[to][from] += 1.0 / pathLength; // 对称更新
        }
        // 回到起点
        int lastCity = antPaths[ant].back();
        int firstCity = antPaths[ant].front();
        pheromoneMatrix[lastCity][firstCity] += 1.0 / pathLength;
        pheromoneMatrix[firstCity][lastCity] += 1.0 / pathLength;
    }
}

// 蚂蚁构造一个解决方案（路径）
QVector<int> AntColonyOptimization::constructSolution(int antIndex) {
    QVector<int> path(numberOfCities);
    QList<bool> visited(numberOfCities, false); // 标记城市是否已访问

    int startCity = rand() % numberOfCities; // 随机选择起始城市
    path[0] = startCity;
    visited[startCity] = true;

    for (int step = 1; step < numberOfCities; ++step) {
        int currentCity = path[step - 1];
        path[step] = selectNextCity(currentCity, visited); // 选择下一个城市
        visited[path[step]] = true;
    }

    return path;
}

// 选择下一个城市
int AntColonyOptimization::selectNextCity(int currentCity, const QVector<bool>& visited) {
    std::vector<double> probabilities(numberOfCities, 0.0);
    double sum = 0.0;

    for (int city = 0; city < numberOfCities; ++city) {
        if (!visited[city]) {
            double pheromone = pow(pheromoneMatrix[currentCity][city], alpha); // 信息素强度的α次幂
            double distance = pow(1.0 / (distanceMatrix[currentCity][city] + 1e-10), beta); // 启发式值的β次幂
            probabilities[city] = pheromone * distance;
            sum += probabilities[city];
        }
    }

    // 轮盘赌选择
    double r = static_cast<double>(rand()) / RAND_MAX * sum;
    double total = 0.0;
    for (int city = 0; city < numberOfCities; ++city) {
        if (!visited[city]) {
            total += probabilities[city];
            if (total >= r) {
                return city;
            }
        }
    }

    // 若轮盘赌未选中，则返回最后一个城市
    return -1;
}

// 蚁群算法求解
QVector<int> AntColonyOptimization::solve() {
    initializePheromones(); // 初始化信息素矩阵

    QVector<int> bestPath;
    double bestLength = INF;

    for (int generation = 0; generation < generations; ++generation) {
        antPaths = std::vector<QVector<int>>(numAnts);
        antPathLengths = std::vector<double>(numAnts, 0.0);

        // 每只蚂蚁构造路径
        for (int ant = 0; ant < numAnts; ++ant) {
            antPaths[ant] = constructSolution(ant);
            antPathLengths[ant] = calculateDistance(antPaths[ant]);

            // 更新最优路径
            if (antPathLengths[ant] < bestLength) {
                bestLength = antPathLengths[ant];
                bestPath = antPaths[ant];
            }
        }

        updatePheromones(); // 更新信息素

        // 输出每一代的最优结果
        qDebug() << "Generation" << generation << "Best distance:" << bestLength;
    }

    return bestPath; // 返回最终最优个体
}
